import cv2
import os

recognizer = cv2.face.LBPHFaceRecognizer_create()
recognizer.read("./data/test.yml")


def get_idx_names():
    files = os.listdir("./users")
    idx_names = {}
    for file in files:
        id = file.split(".")[0]
        name = file.split(".")[1]
        idx_names[id] = name
    return idx_names


#
# idx_names = get_idx_names()
# print(idx_names)


# 检测人脸
def face_detection(img):
    classifier = cv2.CascadeClassifier("./haarcascades/haarcascade_frontalface_alt2.xml")
    faces = classifier.detectMultiScale(
        image=img,
        scaleFactor=1.2,
        minNeighbors=5,
        flags=0,
        minSize=[50, 50],
        maxSize=[200, 200]
    )

    for face in faces:
        x, y, w, h = face
        # 检测人脸
        gray = cv2.cvtColor(img[y:y + h, x:x + w], cv2.COLOR_BGR2GRAY)
        id, confidence = recognizer.predict(gray)
        idx_names = get_idx_names()
        # 使用预测的id与得分判断是否是对应的人
        if confidence < 80:
            cv2.putText(img, idx_names[str(id)], (x, y), 0, 1, color=(0, 0, 255), thickness=2)
        else:
            cv2.putText(img, "unknown", (x, y), 0, 1, color=(0, 0, 255), thickness=2)

        cv2.rectangle(img, (x, y), (x + w, y + h), color=(0, 0, 255), thickness=2)


if __name__ == "__main__":
    cap = cv2.VideoCapture(0)
    while True:
        flags, frame = cap.read()
        if flags is False:
            print("摄像头获取图像失败")
            break
        # 检测人脸
        face_detection(frame)
        cv2.imshow("frame", frame)
        if cv2.waitKey(1) == ord(" "):
            break

    cap.release()
    cv2.destroyAllWindows()
